package com.kanq.handler;

import com.kanq.pojo.context.ChatContext;
import com.kanq.service.QdrantService;
import io.qdrant.client.grpc.Points;

import java.util.List;
import java.util.stream.Collectors;

public class VectorSearchHandler implements ChatHandler {
    private final QdrantService qdrantService;
    private final int topK;

    public VectorSearchHandler(QdrantService qdrantService, int topK) {
        this.qdrantService = qdrantService;
        this.topK = topK;
    }

    @Override
    public void handle(ChatContext context) {
        float[] embedding = qdrantService.getEmbedding(context.getRequest().getQuestion());
        List<Points.ScoredPoint> result = qdrantService.searchVector(
                embedding,
                topK,
                context.getChunkIds()
        );

        List<Long> chunkIds = result.stream()
                .map(x -> x.getId().getNum())
                .collect(Collectors.toList());

        context.setChunkIds(chunkIds);
    }
}